﻿using System.IO;
using System.Runtime.InteropServices;
using Dou.GPU;
using MessagePack;
using MessagePack.Resolvers;
using Unity.Mathematics;
using UnityEditor;
using UnityEngine;
using UnityEngine.Rendering;
using Random = Unity.Mathematics.Random;

namespace Learning
{
    /// <summary>
    /// GPU 动画混合测试
    /// </summary>
    public class GPUInstancingSkinMesh_Blend : MonoBehaviour
    {
        /// <summary>
        /// GPU Instancing 一个批次可提交的数量
        /// 注意：这里设置得越大，DC会越少，但是特殊情况下浪费的带宽会越多（比如只提交10个时，一个批次还是会提交512个数据给GPU，但是有效数据只有10个）
        /// </summary>
        private const int BATCH = 512;
        
        private const float Range = 100000f;
        private readonly Bounds BOUNDS = new Bounds(Vector3.zero, new Vector3(Range, Range, Range));
        
        public int instanceCount = 10000;
        public Mesh instanceMesh;
        public Material instanceMaterial;
        public Object instanceAnimData;
        public string[] instanceAnimNames = new string[0];
        
        private AnimInfo[] _animInfoList;
        private int[] _animNameIndexList;
        
        void Start()
        {
            StaticCompositeResolver.Instance.Register(
                BuiltinResolver.Instance,
                AttributeFormatterResolver.Instance,
                MessagePack.Unity.UnityResolver.Instance,
                PrimitiveObjectResolver.Instance,
                MessagePack.Unity.Extension.UnityBlitWithPrimitiveArrayResolver.Instance,
                StandardResolver.Instance
            );
            var options = MessagePackSerializerOptions.Standard;
            options = options.WithCompression(MessagePackCompression.Lz4BlockArray);
            options = options.WithResolver(StaticCompositeResolver.Instance);
            MessagePackSerializer.DefaultOptions = options;
            
            CreateInfo();
        }

        private void CreateInfo()
        {
            var animData = MessagePackSerializer.Deserialize<GPUAnimData>(File.ReadAllBytes(AssetDatabase.GetAssetPath(instanceAnimData)));
            
            _animInfoList = new AnimInfo[instanceCount];
            _animNameIndexList = new int[instanceCount];

            var col = (int) math.ceil(math.sqrt(instanceCount));
            var random = Random.CreateFromIndex(0);
            
            for (int i = 0; i < instanceCount; i++)
            {
                var animInfo = new AnimInfo();
                var pos = new float3((i % col) * 2f, 0, math.floor(i / col) * 2);
                var rotation = quaternion.Euler(0, random.NextFloat(0f, 360f), 0);
                var s = random.NextFloat(0.8f, 1.2f);
                var scale = new float3(s, s, s);
                animInfo.trans = float4x4.TRS(pos, rotation, scale);
                animInfo.animData = animData;
                _animInfoList[i] = animInfo;

                _animNameIndexList[i] = UnityEngine.Random.Range(0, instanceAnimNames.Length);
                animInfo.Play(instanceAnimNames[_animNameIndexList[i]]);
            }
        }

        void Update()
        {
            // 切换动画
            if (Input.GetKeyUp(KeyCode.A))
            {
                for (int i = 0; i < instanceCount; i++)
                {
                    _animNameIndexList[i]++;
                    _animNameIndexList[i] %= instanceAnimNames.Length;
                    _animInfoList[i].CrossFade(instanceAnimNames[_animNameIndexList[i]]);
                }
            }
            
            UpdateInfo();
            
            var batchCount = (int) math.ceil((float) instanceCount / (float) BATCH);
            for (int i = 0; i < batchCount; i++)
            {
                var startIndex = i * BATCH;
                var len = BATCH;
                if (i == batchCount - 1)
                {
                    len = instanceCount - (BATCH * i);
                }

                // 这里应该做成对象池，这里图方便就暴力创建了
                var batchInfo = new BatchInfo();
                
                // 填充数据，数据长度固定才能使用对象池
                var transArr = new float4x4[BATCH];
                for (int j = 0; j < len; j++)
                {
                    var info = _animInfoList[startIndex + j];
                    transArr[j] = info.trans;
                }
                var transBuf = new GraphicsBuffer(GraphicsBuffer.Target.Structured, BATCH, Marshal.SizeOf(typeof(float4x4)));
                transBuf.SetData(transArr);
                
                var animDataArr = new AnimData[BATCH];
                for (int j = 0; j < len; j++)
                {
                    var data = new AnimData();
                    data.playingIndex = _animInfoList[startIndex + j].playingIndex;
                    data.isCrossBlending = _animInfoList[startIndex + j].isCrossBlending ? 1 : 0;
                    data.lastPlayingIndex = _animInfoList[startIndex + j].lastPlayingIndex;
                    data.crossFadeProgress = _animInfoList[startIndex + j].crossFadeProgress;
                    animDataArr[j] = data;
                    
                    Debug.Log($"playingIndex:{data.playingIndex}, isCrossBlending:{data.isCrossBlending}, lastPlayingIndex:{data.lastPlayingIndex}, crossFadeProgress:{data.crossFadeProgress}");
                }
                var animDataBuf = new GraphicsBuffer(GraphicsBuffer.Target.Structured, BATCH, Marshal.SizeOf(typeof(AnimData)));
                animDataBuf.SetData(animDataArr);
                
                batchInfo.block = new MaterialPropertyBlock();
                batchInfo.block.SetFloat("frameWidth", 1f / _animInfoList[0].animData.frameCount);
                batchInfo.block.SetFloat("frameHeight", 1f / (_animInfoList[0].animData.boneCount * 3));
                batchInfo.block.SetBuffer("transBuf", transBuf);
                batchInfo.block.SetBuffer("animDataBuf", animDataBuf);
                
                var args = new uint[5] {0, 0, 0, 0, 0};
                args[0] = instanceMesh.GetIndexCount(0); // 每个实例的索引数
                args[1] = (uint) len; // 实例数
                args[2] = instanceMesh.GetIndexStart(0); // 起始索引位置
                args[3] = instanceMesh.GetBaseVertex(0); // 基顶点位置
                args[4] = 0; // 起始实例位置
                batchInfo.argsBuf = new GraphicsBuffer(GraphicsBuffer.Target.IndirectArguments, 1, 5 * sizeof(uint));
                batchInfo.argsBuf.SetData(args);

                // 绘制
                Graphics.DrawMeshInstancedIndirect(instanceMesh, 0, instanceMaterial, BOUNDS, batchInfo.argsBuf,
                    0, batchInfo.block, ShadowCastingMode.On, true, LayerMask.NameToLayer("Character"));
            }
        }

        private void UpdateInfo()
        {
            for (int i = 0; i < instanceCount; i++)
            {
                _animInfoList[i].Update();
            }
        }
    }
}
